#include "header.h"

jpeg_header jpeg_header_info;
y4m_stream_info_t stream_info;
int pixels_move_per_block=8;
uint8_t *plane_background=NULL;
unsigned long dl;
//after N_frame write 1 image
int N_image=1;
int Tv=4,N=4;
// after "update" frames update back_ground
int update_background =100;
int quality =80;
short *M,*V,*delta;
int num_blocks;
void read_y4m_info(int fd){
   int ret=0;
   ret=y4m_read_stream_header(fd, &stream_info);
   if(ret != Y4M_OK) y4m_strerr(ret);
   jpeg_header_info.width=y4m_si_get_width(&stream_info);
   jpeg_header_info.height =y4m_si_get_height(&stream_info);
   jpeg_header_info.num_components=1;
   jpeg_header_info.color_space=JCS_GRAYSCALE;
   
}
int get_size(char *filename){
   FILE *fp=fopen(filename,"rb");
   if(fp==NULL){
     fprintf(stderr,"Cannot open file to get size");
     exit(1);
   }
   fseek(fp,0,SEEK_END);
   int x=ftell(fp);
   return x;
}
void block_rebuild(int stt_block,uint8_t *plane){
     int width=jpeg_header_info.width;
     int block_per_width=width/BLOCK_SIZE; 
     unsigned long start_point=(stt_block/block_per_width)*BLOCK_SIZE*BLOCK_SIZE*block_per_width + (stt_block%block_per_width)*BLOCK_SIZE; 
     int i,j,location;
     int num_pixels_move=0;
     for(i=0;i<BLOCK_SIZE;i++)
       for(j=0;j<BLOCK_SIZE;j++){
           location=start_point+i*width+j;
           if(plane[location] < (PIXEL_MAX-5))
               num_pixels_move++;           
     }
     if(num_pixels_move < pixels_move_per_block){
        for(i=0;i<BLOCK_SIZE;i++)
          for(j=0;j<BLOCK_SIZE;j++){
             location=start_point+i*width+j;
             plane[location]=plane_background[location];
          }         
     }

}


int main(int argc, char* argv[])
{
    int opt,ret=0,frame_count=0,size;
    char *infile=NULL,*name=NULL;
    char name2[20];
    strcpy(name2,"output"); name=name2;
    unsigned long i,sum_size_jpeg=0.0;
    char outfile[20],rebuild[20];
    float mean_ssim=0.0,mean_psnr=0.0;
    float psnr,ssim;

    while ((opt = getopt(argc, argv, "i:q:p:t:n:u:o:d:")) != -1) {
         switch (opt) {
              case 'o':
                   name=optarg;
                   break;
              case 'q':
                  quality=atoi(optarg); 
                  break;
              case 'p':
                   pixels_move_per_block = atoi(optarg);
                   break;
             case 'i':
                  infile=optarg;
                  break;
             case 't':
                  Tv=atoi(optarg);
                  break;
             case 'n':
                  N=atoi(optarg);
                  break;
             case 'd':
                  N_image=atoi(optarg);
                  break;
             case 'u':
                  update_background =atoi(optarg);
                  break;
             default:
                break;
                                    
       }
    }
    if(infile==NULL){
        printf("You have to enter infile name after option i \n");
        exit(1);
    }	

    if (optind > argc) {
           fprintf(stderr, "Expected argument after options\n");
           exit(EXIT_FAILURE);
    }

    printf("Quality: %i\n",quality);
    printf("Pixel_move_per_block: %i\n",pixels_move_per_block );
    printf("Tv: %i\n", Tv);
    printf("N: %i\n",N);
    printf("update background %i\n",update_background);
    //Open and read file y4m
    int fd=open(infile,O_RDONLY);
    if(fd==-1){
        perror("Cannot open infile to read");
        exit(1);
    }    
    read_y4m_info(fd);
    //ALLOCATE MEMORY TO STORE DATA
    int np=y4m_si_get_plane_count(&stream_info);
    uint8_t *plane[np];
    unsigned long length_comp[np];
    dl = y4m_si_get_plane_length(&stream_info,0);
    unsigned long total_comp=0;
    for ( i = 0; i < np; i++){
       length_comp[i]=y4m_si_get_plane_length(&stream_info,i);
       total_comp=total_comp+length_comp[i];
       plane[i] = malloc(sizeof(uint8_t)*length_comp[i]);
      // printf("Component %lu size: %lu\n",i, length_comp[i]);
    }
    // READ FRAME DATA
    y4m_frame_info_t frame_info;
    uint8_t *plane_out, *plane_rebuild;
    plane_out=(uint8_t*)malloc(sizeof(uint8_t)*dl);
//    plane_rebuild=(uint8_t*)malloc(sizeof(uint8_t)*dl);
    plane_background=(uint8_t*)malloc(sizeof(uint8_t)*dl);
    M=(short*)malloc(sizeof(short)*dl);
    delta=(short*)malloc(sizeof(short)*dl);
    V=(short*)malloc(sizeof(short)*dl);
    int stt;
    num_blocks=dl/(BLOCK_SIZE*BLOCK_SIZE);
    while(y4m_read_frame_header(fd,&stream_info,&frame_info)==Y4M_OK){

     
        if ( y4m_read_frame_data(fd, &stream_info, &frame_info,plane) == Y4M_OK ){
            sprintf(outfile,"%s%d.jpg",name,frame_count);
            sprintf(rebuild,"rebuild%d.jpg",frame_count);
            //Motion detection using zipfian to write out_file image           
            zipfian(plane[Y_COMP],frame_count,plane_out); 
            write_jpeg_file(outfile,quality,plane_out);
            //Read file out + change to get plane_rebuild.
            plane_rebuild=read_jpeg_file(outfile);
            for(stt=0;stt<num_blocks;stt++){
                block_rebuild(stt,plane_rebuild);
            }
            //Write plane_rebuild
            if((frame_count)%N_image==0 && frame_count >0){
                 write_jpeg_file(rebuild,quality,plane_rebuild);                  

            }           
            size=get_size(outfile);
            sum_size_jpeg=sum_size_jpeg+size;
            psnr=fpsnr(plane[Y_COMP],plane_rebuild);
            mean_psnr=mean_psnr+psnr;
            if(frame_count % update_background==0) ssim=1;
            else 
               ssim=fssim(plane[Y_COMP],plane_rebuild);
            mean_ssim=mean_ssim+ssim;
       }
       frame_count++;
    }
    printf("Number frame: %i\n",frame_count);
    printf("PSRN, SSIM : %f\t%f\n",mean_psnr/frame_count,mean_ssim/frame_count);
    unsigned long size_y4m=get_size(infile);
    float temp=(float)length_comp[0]/total_comp;
    float compress_ratio=(float)size_y4m*temp/(sum_size_jpeg);
    printf("Compression ration: %f\n",compress_ratio);
    //Write result to .txt file
     FILE *ftext=fopen("text.txt","a");
    if(ftext==NULL){
        fprintf(stderr,"cannot open file");
        exit(1);
    }
    fprintf(ftext,"%s\t",infile);
    fprintf(ftext,"%i\t",quality);
    fprintf(ftext,"%i\t%i\t%i\t",pixels_move_per_block,N,Tv);
    fprintf(ftext,"%i\t",update_background);
    fprintf(ftext,"%0.4f\t\t",mean_psnr/frame_count);
    fprintf(ftext,"%0.4f\t\t",mean_ssim/frame_count);
    fprintf(ftext,"%0.4f\n",compress_ratio);
    fclose(ftext);
    //Free memory
    free(plane_background);free(plane_out);free(plane_rebuild);free(M);free(V),free(delta);
    for(i=0;i<np;i++)
        free(plane[i]);
    ret = close(fd);
    if (ret != 0)
       perror("close");
    return 0;
}

